﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace HQHS.Library.Common.Threading
{
    [DebuggerNonUserCode]
    public class KeyedReaderWriterLockSlim<TKey> : IDisposable
    {
        private bool _disposed;
        private readonly ReaderWriterLockSlim _listLock = new ReaderWriterLockSlim();
        private readonly Dictionary<TKey, LockWrapper> _lockList = new Dictionary<TKey, LockWrapper>();
        private readonly LockRecursionPolicy _recursionPolicy;

        #region Constructors
        public KeyedReaderWriterLockSlim()
        {
            this._recursionPolicy = LockRecursionPolicy.NoRecursion;
        }
        public KeyedReaderWriterLockSlim(LockRecursionPolicy recursionPolicy)
        {
            this._recursionPolicy = recursionPolicy;
        }
        #endregion

        #region Public Methods
        public void EnterReadLock(TKey key)
        {
            this.GetLock(key, LockAction.Enter).Lock.EnterReadLock();
        }

        public void EnterUpgradeableReadLock(TKey key)
        {
            this.GetLock(key, LockAction.Enter).Lock.EnterUpgradeableReadLock();
        }

        public void EnterWriteLock(TKey key)
        {
            this.GetLock(key, LockAction.Enter).Lock.EnterWriteLock();
        }

        public void ExitReadLock(TKey key)
        {
            this.GetLock(key, LockAction.Exit).Lock.ExitReadLock();
        }

        public void ExitUpgradeableReadLock(TKey key)
        {
            this.GetLock(key, LockAction.Exit).Lock.ExitUpgradeableReadLock();
        }

        public void ExitWriteLock(TKey key)
        {
            this.GetLock(key, LockAction.Exit).Lock.ExitWriteLock();
        }

        public bool TryEnterReadLock(TKey key, int millisecondsTimeout)
        {
            LockWrapper lw = null;
            if (!this.TryGetLock(key, LockAction.Enter, millisecondsTimeout, out lw))
            {
                return false;
            }
            if (!lw.Lock.TryEnterReadLock(millisecondsTimeout))
            {
                lw.AddLock(Thread.CurrentThread.ManagedThreadId, -1);
                return false;
            }
            return true;
        }

        public bool TryEnterUpgradeableReadLock(TKey key, int millisecondsTimeout)
        {
            LockWrapper lw = null;
            if (!this.TryGetLock(key, LockAction.Enter, millisecondsTimeout, out lw))
            {
                return false;
            }
            if (!lw.Lock.TryEnterUpgradeableReadLock(millisecondsTimeout))
            {
                lw.AddLock(Thread.CurrentThread.ManagedThreadId, -1);
                return false;
            }
            return true;
        }

        public bool TryEnterWriteLock(TKey key, int millisecondsTimeout)
        {
            LockWrapper lw = null;
            if (!this.TryGetLock(key, LockAction.Enter, millisecondsTimeout, out lw))
            {
                return false;
            }
            if (!lw.Lock.TryEnterWriteLock(millisecondsTimeout))
            {
                lw.AddLock(Thread.CurrentThread.ManagedThreadId, -1);
                return false;
            }
            return true;
        }
        #endregion
        #region Non-Public Methods
        private LockWrapper GetLock(TKey key, LockAction change)
        {
            LockWrapper lw = null;
            if (!this.TryGetLock(key, change, -1, out lw))
            {
                throw new ApplicationException();
            }
            return lw;
        }
        private bool TryGetLock(TKey key, LockAction change, int millisecondsTimeout, out LockWrapper lw)
        {
            lw = null;
            if (!this._listLock.TryEnterWriteLock(millisecondsTimeout))
            {
                return false;
            }
            try
            {
                if (!this._lockList.TryGetValue(key, out lw))
                {
                    lw = new LockWrapper(this._recursionPolicy);
                    this._lockList[key] = lw;
                }
                int managedThreadId = Thread.CurrentThread.ManagedThreadId;
                if (change == LockAction.Enter)
                {
                    lw.AddLock(managedThreadId, 1);
                }
                else if (lw.AddLock(managedThreadId, -1) == 0)
                {
                    this._lockList.Remove(key);
                }
            }
            finally
            {
                this._listLock.ExitWriteLock();
            }
            return true;
        }
        #endregion

        #region IDisposable Support
        public void Dispose()
        {
            this.Dispose(true);
            GC.SuppressFinalize(this);
        }
        protected virtual void Dispose(bool disposing)
        {
            if (!this._disposed && disposing)
            {
                if (this._lockList.Count > 0)
                {
                    this._listLock.EnterWriteLock();
                    try
                    {
                        foreach (var wrapper in this._lockList.Values)
                        {
                            wrapper.Lock.Dispose();
                        }
                        this._lockList.Clear();
                    }
                    finally
                    {
                        this._listLock.ExitWriteLock();
                    }
                }
                this._listLock.Dispose();
            }
            this._disposed = true;
        }
        #endregion

        #region Nested Classes
        private class LockWrapper
        {

            private static readonly object _locker = new object();
            private Dictionary<int, int> _threads;

            public readonly ReaderWriterLockSlim Lock;

            #region Constructors

            public LockWrapper(LockRecursionPolicy recursionPolicy)
            {
                this.Lock = new ReaderWriterLockSlim(recursionPolicy);
            }

            #endregion

            #region Methods

            public int AddLock(int threadId, int count)
            {

                int currentCount = default(int);

                lock (_locker)
                {
                    if (this._threads.TryGetValue(threadId, out currentCount))
                    {
                        count += currentCount;

                        if (count == 0)
                        {
                            // Remove threadId if lock count is zero
                            this._threads.Remove(threadId);
                        }
                        else if (count > 0)
                        {
                            //Update lock count
                            this._threads[threadId] = count;
                        }
                        else
                        {
                            throw new InvalidOperationException();
                        }
                    }
                    else
                    {
                        if (count > 0)
                        {
                            this._threads[threadId] = count;
                        }
                        else
                        {
                            throw new InvalidOperationException();
                        }
                    }

                    return _threads.Count;
                }
            }

            #endregion
        }

        #endregion
        private enum LockAction
        {
            Enter = 1,
            Exit = 2
        }
    }
}
